"""6.6 【上下文类】Context 大模型sampling客户端开发"""
import asyncio
import json
from contextlib import AsyncExitStack

from mcp import StdioServerParameters, stdio_client, ClientSession, types, CreateMessageResult
from mcp.client.sse import sse_client
from mcp.shared.context import RequestContext
from mcp.types import TextContent
from openai import OpenAI


class MCPClient:



    def __init__(self):
        self.async_exit_stack = AsyncExitStack()
        self.session = None
        self.deepseek = OpenAI(
            api_key="sk-15af4e21f828460683b16ce9e78b2346",
            base_url="https://api.deepseek.com"
        )

    async def sampling_handler(self,ctx: RequestContext,params: types.CreateMessageRequestParams):
       print("===ctx:",ctx)
       print("===params:",params)
       #一、组装messages,大模型参数
       messages = [
           {
               "role":message.role,
               "content":message.content.text
           } for message in params.messages
       ]
       #二、调用大模型 deepseek
       deepseek_response = self.deepseek.chat.completions.create(
           model="deepseek-chat",
           messages=messages
       )
       # 打印deepseek_response
       print("deepseek_response:", deepseek_response)
       choice_result = deepseek_response.choices[0]
       #四、组装返回的结果
       return CreateMessageResult(
           role="assistant",
           content=TextContent(text=choice_result.message.content, type="text"),
           model="deepseek-chat"
       )







    async def connect_to_server(self,url:str):
        # 一 、创建stdio_client
        client = sse_client(url=url)
        transport = await self.async_exit_stack.enter_async_context(client)
        read_stream, write_stream = transport
        # 二、创建会话client
        client_session = ClientSession(
             read_stream
            ,write_stream
            ,sampling_callback = self.sampling_handler)
        self.session = await self.async_exit_stack.enter_async_context(client_session)
        # 三、初始化会话
        await self.session.initialize()
        # 一、获取server服务端中的工具列表
        response = await self.session.list_tools()
        list_tools = response.tools
        print("【客户端】打印出获取的工具列表：", list_tools)
        arguments = {
            "batch_list": ["1_dt", "2_dt", "3_dt", "4_dt", "5_dt"]
        }

        tool = list_tools[0]
        function_name = tool.name
        # 五、调用工具链
        tool_result = await self.session.call_tool(
            name=function_name,
            arguments=arguments
        )
        print("tool_result:",tool_result)




    #关闭资源
    async def aclose(self):
        await self.async_exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server("http://localhost:8000/sse")
    except Exception as e:
        print(f"连接失败: {e}")
        return
    finally:
        await client.aclose()

if __name__ == "__main__":
    asyncio.run(main())